class WarmupScheduler:
    def __init__(self, decay_epoch, decay_exp, decay_steps):
        self.decay_epoch = [1] + decay_epoch
        self.decay_exp = decay_exp
        self.decay_steps = decay_steps
        self.epoch_index = 0
        self.step_index = decay_steps
        self.cur_decay = []
        self.lr = 0.

    def schedule(self, epoch, optimizer):
        if self.epoch_index < len(self.decay_epoch):
            if epoch == self.decay_epoch[self.epoch_index]:
                if self.step_index == self.decay_steps:
                    self.step_index = 0
                    self.cur_decay.clear()
                    for pg in optimizer.param_groups:
                        lr = pg['lr']
                        if epoch == 1:
                            self.cur_decay.append(lr/self.decay_steps)
                            pg['lr'] = 0
                        else:
                            self.cur_decay.append((lr*self.decay_exp-lr)/self.decay_steps)
                    self.epoch_index += 1
        if self.step_index < self.decay_steps:
            self.step_index += 1
            for i, pg in enumerate(optimizer.param_groups):
                self.lr += self.cur_decay[i]
                pg['lr'] = pg['lr'] + self.cur_decay[i]


class NoScheduler:
    def __init__(self):
        pass

    def schedule(self, epoch, optimizer):
        pass


class SchedulerFactory:
    schedulers = {'WarmupScheduler': WarmupScheduler, 'NoScheduler': NoScheduler}

    @staticmethod
    def build_optimizer(optimizer_cfg):
        cfg = optimizer_cfg.copy()
        scheduler_cls = SchedulerFactory.schedulers[cfg.pop('type')]
        return scheduler_cls(**cfg)


def build_scheduler(scheduler_cfg):
    return SchedulerFactory.build_optimizer(scheduler_cfg)